### pkdgrav_ss.py

"""Classes and functions for accessing PKDGRAV data."""

import struct
import numpy as npy
import os

class ssHeader:
    """
    PKDGRAV ss file header class
    
    t - time (year/(2*pi))
    N - number of particles
    MagicNumber - ss magic number, indicates normal or reduced format [-1]
    """
    def __init__(self, t=0, N=0, magic=0):
        self.t = t
        self.N = N
        self.MagicNumber = magic
        
class ss:
    """
    PKDGRVAV ss file class
    
    header - ss file header
    m - particle mass (solar masses)
    r - particle radius (AU)
    x,y,z - particle coordinates (AU)
    vx,vy,vz - particle velocity componenets (AU/year/(2*pi))
    sx,sy,sz - particle spin components [spin is not used]
    org_idx - original particle index
    color - particle color, indicates type
                [3  = planetesimal
                 2  = Jupiter
                 11 = Saturn]    
    (origin - mass origin hsitogram)
    
    
    Methods:
    
    read(<file>,extras=False)
        - read ss file <file>
        - if extras is True, will try to read origin histograms
        """
    
    def __init__(self):
        self.header = ssHeader()
        self.m = npy.zeros(self.header.N)
        self.r = npy.zeros(self.header.N)
        self.x = npy.zeros(self.header.N)
        self.y = npy.zeros(self.header.N)
        self.z = npy.zeros(self.header.N)
        self.vx = npy.zeros(self.header.N)
        self.vy = npy.zeros(self.header.N)
        self.vz = npy.zeros(self.header.N)
        self.sx = npy.zeros(self.header.N)
        self.sy = npy.zeros(self.header.N)
        self.sz = npy.zeros(self.header.N)
        self.color = npy.zeros(self.header.N).astype(int)
        self.org_idx = npy.zeros(self.header.N).astype(int)

    def read(self, fname, extras=False):
        dbytes = 8
        ibytes = 4
        with open(fname, 'rb') as fo:
            self.header.t = struct.unpack('>d', fo.read(dbytes))[0]
            self.header.N = struct.unpack('>i', fo.read(ibytes))[0]
            self.header.MagicNumber = struct.unpack('>i',fo.read(ibytes))[0]

            # initialise arrays
            self.m = npy.zeros(self.header.N)
            self.r = npy.zeros(self.header.N)
            self.x = npy.zeros(self.header.N)
            self.y = npy.zeros(self.header.N)
            self.z = npy.zeros(self.header.N)
            self.vx = npy.zeros(self.header.N)
            self.vy = npy.zeros(self.header.N)
            self.vz = npy.zeros(self.header.N)
            self.sx = npy.zeros(self.header.N)
            self.sy = npy.zeros(self.header.N)
            self.sz = npy.zeros(self.header.N)
            self.color = npy.zeros(self.header.N).astype(int)
            self.org_idx = npy.zeros(self.header.N).astype(int)

            # read data
            for j in range(0, self.header.N):
                self.m[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.r[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.x[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.y[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.z[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.vx[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.vy[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.vz[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.sx[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.sy[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.sz[j] = struct.unpack('>d', fo.read(dbytes))[0]
                self.color[j] = struct.unpack('>i', fo.read(ibytes))[0]
                self.org_idx[j] = struct.unpack('>i', fo.read(ibytes))[0]
        
        
        if extras:  #try to load origin histograms
            if os.path.isfile(fname+'.origin_bins'):
                self.origin = npy.transpose(npy.loadtxt(fname+'.origin_bins',skiprows=1,unpack=True))
     


if __name__ == "__main__":
    import sys
    import numpy as npy
    import matplotlib.pyplot as plt

    file = 'ssic.ss'
    if len(sys.argv) > 1:
        file = sys.argv[1]

    pkdss = ss()
    pkdss.read(file,extras=True)

    print(pkdss.header.N)   # print number particles
    print(pkdss.m[pkdss.color==2])   # print mass of Jupiter
    

    # calculate orbital elements and plot a vs e
    d=npy.sqrt(pkdss.x**2+pkdss.y**2+pkdss.z**2)
    v2=pkdss.vx**2+pkdss.vy**2+pkdss.vz**2
    rv=pkdss.x*pkdss.vx+pkdss.y*pkdss.vy+pkdss.z*pkdss.vz
    hx=pkdss.y*pkdss.vz - pkdss.z*pkdss.vy
    hy=pkdss.z*pkdss.vx - pkdss.x*pkdss.vz
    hz=pkdss.x*pkdss.vy - pkdss.y*pkdss.vx
    hxy=hx**2+hy**2
    h2=hxy+hz**2

    i=npy.arctan2(npy.sqrt(hxy),hz)

    mu = 1.
    E=0.5*v2 - mu/d
    a=npy.where(E<0,-0.5*mu/E,npy.where(E>0,0.5*mu/E,2*mu/v2))
    e=npy.where(E<0,npy.sqrt(1-h2/(mu*a)),npy.where(E>0,npy.sqrt(1+h2/(mu*a)),1.))

    
    plt.scatter(a,e)
    plt.axis([1.,15.,0.,1.])
    plt.xlabel('a (au)')
    plt.ylabel('e')
    plt.show()
           
